from transformers import BertTokenizer
from transformers import BertForSequenceClassification

tokenizer = BertTokenizer.from_pretrained('/gemini/data-1')
model = BertForSequenceClassification.from_pretrained('/gemini/data-1')

tokenizer.add_tokens(['[rationale]', '[\\rationale]', '[class]', '[\\class]',
                      '[query]', '[\\query]', '[answer]', '[\\answer]'])

model.resize_token_embeddings(len(tokenizer))
model.save_pretrained("/gemini/code/RationaleLoop/save")
tokenizer.save_pretrained('/gemini/code/RationaleLoop/save')
print('Done !')